/*
 *   Copyright 2021 zzh
 *
 *   Licensed under the Apache License, Version 2.0 (the "License");
 *   you may not use this file except in compliance with the License.
 *   You may obtain a copy of the License at
 *
 *       http://www.apache.org/licenses/LICENSE-2.0
 *
 *   Unless required by applicable law or agreed to in writing, software
 *   distributed under the License is distributed on an "AS IS" BASIS,
 *   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *   See the License for the specific language governing permissions and
 *   limitations under the License.
 *
 */

package com.aduib.boot.common.util;

import sun.misc.Unsafe;

import java.nio.Buffer;
import java.nio.ByteBuffer;
import java.nio.charset.Charset;

/**
 * @description: ByteBufferUtils
 * @author: zzh
 * @date: 2021/12/2 10:14
 */
public class ByteBufferUtils {

  public static final int DEFAULT_BUFFER_SIZE = 1024;

  static final Unsafe theUnsafe;
  /** The offset to the first element in a byte array. */
  static final long BYTE_ARRAY_BASE_OFFSET;

  static final long DIRECT_BUFFER_ADDRESS_OFFSET;
  // 1M, copied from java.nio.Bits (unfortunately a package-private class)
  private static final long UNSAFE_COPY_THRESHOLD = 1 << 20;
  private static final long MIN_COPY_THRESHOLD = 6;

  static {
    theUnsafe = SystemUtils.UNSAFE;

    try {
      BYTE_ARRAY_BASE_OFFSET = theUnsafe.arrayBaseOffset(byte[].class);
      DIRECT_BUFFER_ADDRESS_OFFSET =
          theUnsafe.objectFieldOffset(Buffer.class.getDeclaredField("address"));
    } catch (Exception e) {
      throw new AssertionError(e);
    }

    // sanity check - this should never fail
    if (theUnsafe.arrayIndexScale(byte[].class) != 1) {
      throw new AssertionError();
    }
  }

  /**
   * 写完数据切换为读模式 写完数据，需要开始读的时候，将postion复位到0，并将limit设为当前postion。
   *
   * @param byteBuffer
   */
  public static void flip(ByteBuffer byteBuffer) {
    if (byteBuffer.position() != 0) {
      byteBuffer.flip();
    }
  }

  /**
   * 为某一读过的位置做标记，便于某些时候回退到该位置。
   *
   * @param byteBuffer
   */
  public static void mark(ByteBuffer byteBuffer) {
    if (byteBuffer.position() != 0) {
      byteBuffer.mark();
    }
  }

  /**
   * 创建一个指定大小的bytebuffer
   *
   * @param size
   * @return
   */
  public static ByteBuffer allocateDirect(int size) {
    return ByteBuffer.allocateDirect(size <= 0 ? DEFAULT_BUFFER_SIZE : size);
  }

  /**
   * 创建一个指定大小的bytebuffer
   *
   * @param size
   * @return
   */
  public static ByteBuffer allocate(int size) {
    return ByteBuffer.allocate(size <= 0 ? DEFAULT_BUFFER_SIZE : size);
  }

  /**
   * 拷贝到一个新的ByteBuffer
   *
   * @param src 源ByteBuffer
   * @param start 起始位置（包括）
   * @param end 结束位置（不包括）
   * @return 新的ByteBuffer
   */
  public static ByteBuffer copy(ByteBuffer src, int start, int end) {
      ByteBuffer buffer = ByteBuffer.allocate(end - start);
      copy(src, buffer);
      return buffer;
  }

  /**
   * 通过调用Buffer.mark()方法，可以标记Buffer中的一个特定position。之后可以通过调用Buffer.reset()方法恢复到这个position。
   *
   * @param byteBuffer
   */
  public static void reset(ByteBuffer byteBuffer) {
    if (byteBuffer.position() != 0) {
      byteBuffer.reset();
    }
  }

  /**
   * 把position设为0，把limit设为capacity，一般在把数据写入Buffer前调用。
   *
   * @param byteBuffer
   */
  public static void clear(ByteBuffer byteBuffer) {
    if (byteBuffer.position() != 0) {
      byteBuffer.clear();
    }
  }

  /**
   * 把position设为0，limit不变，一般在把数据重写入Buffer前调用。
   *
   * @param byteBuffer
   */
  public static void rewind(ByteBuffer byteBuffer) {
    if (byteBuffer.position() != 0) {
      byteBuffer.rewind();
    }
  }

  /**
   * 拷贝ByteBuffer
   *
   * @param src 源ByteBuffer
   * @param dest 目标ByteBuffer
   * @return 目标ByteBuffer
   */
  public static void copy(ByteBuffer src, ByteBuffer dest) {
    copy(src, dest, Math.min(src.limit(), dest.remaining()));
  }

  /**
   * 拷贝ByteBuffer
   *
   * @param src 源ByteBuffer
   * @param dest 目标ByteBuffer
   * @param length 长度
   * @return 目标ByteBuffer
   */
  public static void copy(ByteBuffer src, ByteBuffer dest, int length) {
    copy(src, src.position(), dest, dest.position(), length);
  }

  public static void copy(ByteBuffer src, int srcPosition, byte[] trg, int trgPosition, int length) {
    if (src.hasArray()) {
      System.arraycopy(src.array(), src.arrayOffset() + srcPosition, trg, trgPosition, length);
    } else {
      copy(
          null,
          srcPosition + theUnsafe.getLong(src, DIRECT_BUFFER_ADDRESS_OFFSET),
          trg,
          trgPosition,
          length);
    }
  }

  public static void copy(
      ByteBuffer srcBuf, int srcPosition, ByteBuffer trgBuf, int trgPosition, int length) {
    Object src;
    long srcOffset;
    if (srcBuf.hasArray()) {
      src = srcBuf.array();
      srcOffset = BYTE_ARRAY_BASE_OFFSET + srcBuf.arrayOffset();
    } else {
      src = null;
      srcOffset = theUnsafe.getLong(srcBuf, DIRECT_BUFFER_ADDRESS_OFFSET);
    }
    copy(src, srcOffset + srcPosition, trgBuf, trgPosition, length);
  }

  public static void copy(
      Object src, long srcOffset, ByteBuffer trgBuf, int trgPosition, int length) {
    if (trgBuf.hasArray()) {
      copy(src, srcOffset, trgBuf.array(), trgBuf.arrayOffset() + trgPosition, length);
    } else {
      copy(
          src,
          srcOffset,
          null,
          trgPosition + theUnsafe.getLong(trgBuf, DIRECT_BUFFER_ADDRESS_OFFSET),
          length);
    }
  }

  public static void copy(Object src, long srcOffset, byte[] trg, int trgPosition, int length) {
    if (length <= MIN_COPY_THRESHOLD) {
      for (int i = 0; i < length; i++) {
        trg[trgPosition + i] = theUnsafe.getByte(src, srcOffset + i);
      }
    } else {
      copy(src, srcOffset, trg, BYTE_ARRAY_BASE_OFFSET + trgPosition, length);
    }
  }

  public static void copy(Object src, long srcOffset, Object dst, long dstOffset, long length) {
    while (length > 0) {
      long size = Math.min(length, UNSAFE_COPY_THRESHOLD);
      // if src or dst are null, the offsets are absolute base addresses:
      theUnsafe.copyMemory(src, srcOffset, dst, dstOffset, size);
      length -= size;
      srcOffset += size;
      dstOffset += size;
    }
  }

  /**
   * 读取剩余部分并转为UTF-8编码字符串
   *
   * @param buffer ByteBuffer
   * @return 字符串
   */
  public static String readUtf8Str(ByteBuffer buffer) {
    return readStr(buffer, CharsetUtils.CHARSET_UTF_8);
  }

  /**
   * 读取剩余部分并转为字符串
   *
   * @param buffer ByteBuffer
   * @param charset 编码
   * @return 字符串
   */
  public static String readStr(ByteBuffer buffer, Charset charset) {
    return StrUtils.str(readBytes(buffer), charset);
  }

  /**
   * 读取剩余部分bytes<br>
   *
   * @param buffer ByteBuffer
   * @return bytes
   */
  public static byte[] readBytes(ByteBuffer buffer) {
    final int remaining = buffer.remaining();
    byte[] ab = new byte[remaining];
    buffer.get(ab);
    return ab;
  }

  /**
   * 读取指定长度的bytes<br>
   * 如果长度不足，则读取剩余部分，此时buffer必须为读模式
   *
   * @param buffer ByteBuffer
   * @param maxLength 最大长度
   * @return bytes
   */
  public static byte[] readBytes(ByteBuffer buffer, int maxLength) {
    final int remaining = buffer.remaining();
    if (maxLength > remaining) {
      maxLength = remaining;
    }
    byte[] ab = new byte[maxLength];
    buffer.get(ab);
    return ab;
  }

  /**
   * 读取指定区间的数据
   *
   * @param buffer {@link ByteBuffer}
   * @param start 开始位置
   * @param end 结束位置
   * @return bytes
   */
  public static byte[] readBytes(ByteBuffer buffer, int start, int end) {
    byte[] bs = new byte[end - start];
    System.arraycopy(buffer.array(), start, bs, 0, bs.length);
    return bs;
  }

  /**
   * 一行的末尾位置，查找位置时位移ByteBuffer到结束位置
   *
   * @param buffer {@link ByteBuffer}
   * @return 末尾位置，未找到或达到最大长度返回-1
   */
  public static int lineEnd(ByteBuffer buffer) {
    return lineEnd(buffer, buffer.remaining());
  }

  /**
   * 一行的末尾位置，查找位置时位移ByteBuffer到结束位置<br>
   * 支持的换行符如下：
   *
   * <pre>
   * 1. \r\n
   * 2. \n
   * </pre>
   *
   * @param buffer {@link ByteBuffer}
   * @param maxLength 读取最大长度
   * @return 末尾位置，未找到或达到最大长度返回-1
   */
  public static int lineEnd(ByteBuffer buffer, int maxLength) {
    int primitivePosition = buffer.position();
    boolean canEnd = false;
    int charIndex = primitivePosition;
    byte b;
    while (buffer.hasRemaining()) {
      b = buffer.get();
      charIndex++;
      if (b == StrUtils.C_CR) {
        canEnd = true;
      } else if (b == StrUtils.C_LF) {
        return canEnd ? charIndex - 2 : charIndex - 1;
      } else {
        // 只有\r无法确认换行
        canEnd = false;
      }

      if (charIndex - primitivePosition > maxLength) {
        // 查找到尽头，未找到，还原位置
        buffer.position(primitivePosition);
        throw new IndexOutOfBoundsException(
            StrUtils.format("Position is out of maxLength: {}", maxLength));
      }
    }

    // 查找到buffer尽头，未找到，还原位置
    buffer.position(primitivePosition);
    // 读到结束位置
    return -1;
  }

  /**
   * 读取一行，如果buffer中最后一部分并非完整一行，则返回null<br>
   * 支持的换行符如下：
   *
   * <pre>
   * 1. \r\n
   * 2. \n
   * </pre>
   *
   * @param buffer ByteBuffer
   * @param charset 编码
   * @return 一行
   */
  public static String readLine(ByteBuffer buffer, Charset charset) {
    final int startPosition = buffer.position();
    final int endPosition = lineEnd(buffer);

    if (endPosition > startPosition) {
      byte[] bs = readBytes(buffer, startPosition, endPosition);
      return StrUtils.str(bs, charset);
    } else if (endPosition == startPosition) {
      return StrUtils.EMPTY;
    }

    return null;
  }
}
